import cv2
import numpy as np
import open3d as o3d


class StereoReconstructor:
    """
    双目三维重建核心类
    输入：重投影矩阵 Q、校正映射参数、左右图像
    输出：三维点云
    """

    def __init__(self, Q, left_map_x, left_map_y, right_map_x, right_map_y, roi_left, roi_right):
        self.Q = Q
        self.left_map_x = left_map_x
        self.left_map_y = left_map_y
        self.right_map_x = right_map_x
        self.right_map_y = right_map_y
        self.roi_left = roi_left  # (x,y,w,h)
        self.roi_right = roi_right

    def rectify_images(self, img_left, img_right):
        """立体校正并裁剪有效区域"""
        rect_left = cv2.remap(img_left, self.left_map_x, self.left_map_y, cv2.INTER_LINEAR)
        rect_right = cv2.remap(img_right, self.right_map_x, self.right_map_y, cv2.INTER_LINEAR)
        # 根据 roi 裁剪各自图像
        x1, y1, w1, h1 = self.roi_left
        x2, y2, w2, h2 = self.roi_right

        # 计算公共 ROI（交集）
        x_common = max(x1, x2)
        y_common = max(y1, y2)
        w_common = min(x1 + w1, x2 + w2) - x_common
        h_common = min(y1 + h1, y2 + h2) - y_common

        rect_left = rect_left[y_common:y_common + h_common, x_common:x_common + w_common]
        rect_right = rect_right[y_common:y_common + h_common, x_common:x_common + w_common]
        return rect_left, rect_right

    def compute_disparity(self, img_left, img_right, use_filter=True):
        gray_left = cv2.cvtColor(img_left, cv2.COLOR_BGR2GRAY)
        gray_right = cv2.cvtColor(img_right, cv2.COLOR_BGR2GRAY)

        # SGBM参数配置
        window_size = 5
        min_disp = 0
        num_disp = 128 - min_disp  # 视差范围

        stereo_left = cv2.StereoSGBM_create(
            minDisparity=min_disp,
            numDisparities=num_disp,
            blockSize=window_size,
            P1=8 * 3 * window_size ** 2,
            P2=32 * 3 * window_size ** 2,
            disp12MaxDiff=1,
            uniquenessRatio=15,
            speckleWindowSize=100,
            speckleRange=1
        )
        stereo_right = cv2.ximgproc.createRightMatcher(stereo_left)

        disp_left = stereo_left.compute(gray_left, gray_right).astype(np.float32) / 16.0
        disp_right = stereo_right.compute(gray_right, gray_left).astype(np.float32) / 16.0

        if use_filter:
            wls_filter = cv2.ximgproc.createDisparityWLSFilter(matcher_left=stereo_left)
            disp_left = wls_filter.filter(disp_left, gray_left, None, disp_right)

        # 清理无效值：将 inf、-inf 和 -16 置为0
        disp_left[~np.isfinite(disp_left)] = 0
        disp_left[disp_left < 0] = 0

        return disp_left

    def depth_to_3d(self, disparity):
        """通过视差图计算三维坐标"""
        points_3d = cv2.reprojectImageTo3D(disparity, self.Q)
        return points_3d

    def create_point_cloud(self, points_3d, color_img):
        """生成Open3D点云对象"""
        # 过滤无效点：排除视差值接近最小值的点（可根据需要调整过滤条件）
        mask = disparity > 0  # 只保留视差大于0的点
        pcd = o3d.geometry.PointCloud()
        pcd.points = o3d.utility.Vector3dVector(points_3d[mask])
        pcd.colors = o3d.utility.Vector3dVector(color_img[mask] / 255.0)
        return pcd


def main():
    # ----------------------------
    # 标定参数及立体校正参数（这里假设部分参数已经修正）
    # ----------------------------
    # 1. 定义 Q 矩阵（4x4）
    Q = np.array([
        [1.0, 0.0, 0.0, -7.82334202e+02],
        [0.0, 1.0, 0.0, -3.18832064e+02],
        [0.0, 0.0, 0.0, 2.39754133e+03],
        [0.0, 0.0, 1.11818766e-02, 0.0]
    ], dtype=np.float64)

    # 2. 定义旋转矩阵 R1, R2（3x3）
    R1 = np.array([
        [0.99925563, -0.01089273, 0.03700723],
        [0.01039314, 0.99985261, 0.01366563],
        [-0.03715063, -0.01327084, 0.99922155]
    ], dtype=np.float64)

    R2 = np.array([
        [0.9999786, 0.00571043, 0.00319094],
        [-0.00566693, 0.99989308, -0.01347986],
        [-0.00326757, 0.01346148, 0.99990405]
    ], dtype=np.float64)

    # 3. 定义投影矩阵 P1, P2（3x4）
    P1 = np.array([
        [2.39754133e+03, 0.0, 7.82334202e+02, 0.0],
        [0.0, 2.39754133e+03, 3.18832064e+02, 0.0],
        [0.0, 0.0, 1.0, 0.0]
    ], dtype=np.float64)

    P2 = np.array([
        [2.39754133e+03, 0.0, 7.82334202e+02, -2.14413145e+05],
        [0.0, 2.39754133e+03, 3.18832064e+02, 0.0],
        [0.0, 0.0, 1.0, 0.0]
    ], dtype=np.float64)

    # 4. 畸变系数（假设无畸变）
    D1 = np.zeros((4, 1), dtype=np.float64)
    D2 = np.zeros((4, 1), dtype=np.float64)

    # 5. 读取图像并获取尺寸
    img_left = cv2.imread("data1/left_1.jpg")
    img_right = cv2.imread("data1/right_1.jpg")
    if img_left is None or img_right is None:
        raise IOError("读取图像失败，请检查文件路径")
    h, w = img_left.shape[:2]
    image_size = (w, h)

    # ----------------------------
    # 立体校正参数计算（alpha 设置为0.5）
    # ----------------------------
    ret = cv2.stereoRectify(
        mtx_main, dist_main, mtx_side, dist_side,
        image_size, R, T,
        flags=cv2.CALIB_ZERO_DISPARITY,
        alpha=0.5
    )
    # stereoRectify 返回 (R1,R2,P1,P2,Q,roi1,roi2)
    R1_calib, R2_calib, P1_calib, P2_calib, Q, roi1, roi2 = ret

    # 6. 计算校正映射（注意这里内参采用投影矩阵P1、P2的前三列）
    K1_calib = P1_calib[:3, :3]
    K2_calib = P2_calib[:3, :3]
    left_map_x, left_map_y = cv2.initUndistortRectifyMap(
        mtx_main, D1, R1_calib, K1_calib, image_size, cv2.CV_32FC1
    )
    right_map_x, right_map_y = cv2.initUndistortRectifyMap(
        mtx_side, D2, R2_calib, K2_calib, image_size, cv2.CV_32FC1
    )

    # ----------------------------
    # 重建流程
    # ----------------------------
    # 构造重建器，将 roi 一并传入以便裁剪校正后图像
    reconstructor = StereoReconstructor(Q, left_map_x, left_map_y, right_map_x, right_map_y, roi1, roi2)

    # 对输入图像进行立体校正和裁剪
    rect_left, rect_right = reconstructor.rectify_images(img_left, img_right)

    # 计算视差图
    global disparity  # 用于 create_point_cloud 中过滤
    disparity = reconstructor.compute_disparity(rect_left, rect_right)

    # 打印视差图统计信息
    print("视差图最小值：", np.nanmin(disparity))
    print("视差图最大值：", np.nanmax(disparity))
    print("有效点数量：", np.count_nonzero(disparity > (np.nanmin(disparity) + 1e-3)))

    # 生成三维点云
    points_3d = reconstructor.depth_to_3d(disparity)

    # 创建并保存点云
    pcd = reconstructor.create_point_cloud(points_3d, rect_left)
    o3d.io.write_point_cloud("output1.ply", pcd)

    # 可视化点云
    o3d.visualization.draw_geometries([pcd])

    # 显示校正后的图像
    cv2.imshow("rectified left", rect_left)
    cv2.imshow("rectified right", rect_right)
    cv2.waitKey(0)
    cv2.destroyAllWindows()


if __name__ == "__main__":
    # 定义标定内参及畸变参数（与上面的示例保持一致）
    mtx_main = np.array([[1093.23802, 0, 635.746975],
                         [0, 1096.467663, 337.6218325],
                         [0, 0, 1]], dtype=np.float64)
    dist_main = np.array([-0.149254721, 0.54330647, 0, 0, 0], dtype=np.float64).reshape(1, -1)
    mtx_side = np.array([[1091.510188, 0, 659.0351622],
                         [0, 1093.072199, 298.8419753],
                         [0, 0, 1]], dtype=np.float64)
    dist_side = np.array([-0.034262054, -0.123556241, 0, 0, 0], dtype=np.float64).reshape(1, -1)
    frame_size = (1280, 720)
    R = np.array([[0.999296745, -0.016515227, 0.033663969],
                  [0.0155981, 0.999504858, 0.027326502],
                  [-0.034098604, -0.026782191, 0.999059558]], dtype=np.float64)
    T = np.array([-89.42851336, -0.510686004, -0.285366743], dtype=np.float64).reshape(3, 1)

    main()
